function [HMMresults,statemaps] = osl_hmm_groupinference(data_files,hmmdir,todo,options)
% Runs a group HMM analysis on amplitude envelopes of source reconstructed
% MEG data using the OAT beamformer OR any arbitrary data saved in .mat
% data files.
%
% osl_hmm_groupinference(data_files,hmmdir,todo,env_settings,hmm_settings)
%
% INPUTS:
%
% data_files - cell array containing list of OAT source recon results 
%              (oat_stage_results) containing continuous data OR .mat files 
%              containing data [observations x timepoints]
%  
% hmmdir     - directory within which to save HMM results
%
% todo       - binary flags for each stage of the analysis [0/1]:
%              .envelope - computes amplitude envelopes from source
%                        reconstructed data
%              .concat   - concatenates data and performs PCA dimensionality
%                        reduction
%              .hmm      - infers the HMM
%              .output   - output the state maps
%
%
% options    - struct with settings for envelope, concatenation and hmm stages
%              with the following fields:
%              .envelope - settings for envelope computation with fields:
%                          .fsample    - specify sampling rate if NOT using 
%                                        OAT source recon results
%                          .windowsize - window size for moving average filter [s]
%                            (default 0.1)
%                          .overlap    - overlap for moving average filter [0-1]
%                            (default 0.75)
%                          .filenames  - filenames to save/load envelopes from
%                            (default <hmmdir>/envelopes/<BFfiles{subnum}>.mat)
%             
%              .concat   - settings for temporal concatenation with fields:
%                          .norm_subs  - apply subject variance normalisation [0/1]
%                            (default 1)
%                          .pcadim     - dimensionality to use 
%                            (default 40)                       
%                          .whiten     - apply whitening [0/1]                      
%                            (default 1)
%                          .filename   - filename to save/load concatenated data from
%                            (default <hmmdir>/env_concat.mat)
%
%              .hmm      - settings for HMM inference with fields:                
%                          .nstates    - number of states to infer 
%                            (default 8)
%                          .nreps      - number of repeat inferences to run 
%                            (default 5)
%                          .filename   - filename to save/load HMM results from
%                            (default <hmmdir>/HMM.mat)
%      
%              .output   - settings for the HMM output with fields
%                          .method     - maps/matrix to output:
%                                        'pcorr' - partial correlation maps
%                                        'var'   - variance maps
%                                        'cov'   - covariance matrices
%                            (default 'pcorr')
%                          .filename   - filename to save/load HMM results from
%                            (default <hmmdir>/HMM<_method.*>)
%                       
%
% OUTPUTS:
%
% HMMresults - filename of the inferred HMM results
%
% statemaps  - filename of the HMM state maps
%
%
% AB 2014

%global OSLDIR

HMMresults = [];
statemaps  = [];

data_files = cellstr(data_files);

data_fnames = cell(size(data_files));
for f = 1:numel(data_files)
  [~,data_fnames{f},~] = fileparts(data_files{f});
end

if ~isdir(hmmdir)
  mkdir(hmmdir); 
end


% Set up directories:
if isfield(options.envelope,'filenames') && ~isempty(options.envelope.filenames)
    filenames.envelope = cell(size(options.envelope.filenames));
    [pathstr,filestr] = cellfun(@fileparts,options.envelope.filenames,'uniformoutput',0);
    for f = 1:length(options.envelope.filenames)
        if strcmp(pathstr(f),'')
            pathstr{f} = fullfile(hmmdir,'envelopes');
            if ~isdir(pathstr{f})
                mkdir(pathstr{f})
            end
        end
        filenames.envelope{f} = fullfile(pathstr{f},filestr{f});
    end
else
    pathstr = fullfile(hmmdir,'envelopes');
    if ~isdir(pathstr)
        mkdir(pathstr)
    end
    filenames.envelope = fullfile(pathstr,data_fnames);
end
filenames.envelope = strcat(filenames.envelope, '.mat');

if isfield(options.concat,'filename') && ~isempty(options.concat.filename)
    [pathstr,filestr] = fileparts(options.concat.filename);
        if strcmp(pathstr,'')
            pathstr = hmmdir;
        end
        filenames.concat = fullfile(pathstr,filestr);
else
    filenames.concat = fullfile(hmmdir,'env_concat');
end
filenames.concat = [filenames.concat '.mat'];

if isfield(options.hmm,'filename') && ~isempty(options.hmm.filename)
    [pathstr,filestr] = fileparts(options.hmm.filenam);
        if strcmp(pathstr,'')
            pathstr = hmmdir;
        end
        filenames.hmm = fullfile(pathstr,filestr);
else
    filenames.hmm = fullfile(hmmdir,'hmm');
end
filenames.hmm = [filenames.hmm '.mat'];

if isfield(options.output,'filename') && ~isempty(options.output.filename)
    [pathstr,filestr] = fileparts(options.output.filenam);
        if strcmp(pathstr,'')
            pathstr = hmmdir;
        end
        filenames.output = fullfile(pathstr,filestr);
else
    filenames.output = fullfile(hmmdir,'hmm');
end



% Default todo settings
try todo.envelope = todo.envelope; catch, todo.envelope = 1; end
try todo.concat   = todo.concat;   catch, todo.concat   = 1; end
try todo.hmm      = todo.hmm;      catch, todo.hmm      = 1; end

% Default envelope settings
try windowsize = options.envelope.windowsize; catch, windowsize = 0.1;  end
try overlap    = options.envelope.overlap;    catch, overlap    = 0.75; end

% Default concatenation settings
try norm_subs = options.concat.norm_subs;  catch, norm_subs  = 1;  end
try pcadim    = options.concat.pcadim;     catch, pcadim     = 40; end
try whiten    = options.concat.whiten;     catch, whiten     = 1;  end

% Default HMM settings
try nstates = options.hmm.nstates; catch, nstates = 8; end
try nreps   = options.hmm.nreps;   catch, nreps   = 5; end
    
% Default output settings
try output_method = options.output.method; catch, output_method = 'pcorr'; end



%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                         %
%                    C O M P U T E   E N V E L O P E S                    %
%                                                                         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if todo.envelope
  
  for subnum = 1:length(data_files)
    
    [pathstr,filestr] = fileparts(data_files{subnum});
  
    disp(['Computing envelope data for ' filestr]);
    
    dat = load(data_files{subnum});   
    try 
        dat = dat.oat_stage_results;
        isOSL = 1;
    catch
        if length(fieldnames(dat))>1
            error('.mat data_files should contain only one variable')
        else
            dat = dat.(char(fieldnames(dat)));
        end
        isOSL = 0;
    end
    
    
    if isOSL
        
        dat.source_recon.dirname = pathstr;
        tbad = ~dat.samples2use;
        
        recon = get_source_timecourses(dat,'noreject','norecon');
        
        % Apply hilbert transform in sensor space:
        recon.sensor_data = transpose(hilbert(recon.sensor_data'));
        
        % Preallocation
        Nvoxels  = length(recon.weights);
        fsample  = dat.BF.data.D.fsample;
        winsize  = windowsize * fsample;
        t        = dat.BF.data.D.time;
        t(tbad)  = nan;
    
    else
        Nvoxels = size(dat,1);
        try 
            fsample = options.envelope.fsample;    
        catch  
            error('Must specifiy options.envelop.fsample if NOT using OAT results')
        end
        winsize  = windowsize * fsample;
        t = (1:length(dat))/fsample;
        tbad = false(1,size(dat,2));

    end
        

    if winsize ~= 0
      env = zeros(Nvoxels, length(osl_movavg(t,t,winsize,overlap,0)));
    else
      env = zeros(Nvoxels, length(t));
    end
    
    ft_progress('init', 'etf');
    for vox=1:Nvoxels
        ft_progress(vox/Nvoxels);
        
        if isOSL
            tdata = get_source_timecourses_recon(recon,vox);
            tdata = abs(tdata{1});
        else
            tdata = abs(hilbert(dat(vox,:)));
        end
        
        HE = sqrt(sum(tdata.^2,1)); % Vector beamformer compatible
        HE(:,tbad) = nan;
        
        if winsize ~= 0
            [env(vox,:),t_avg] = osl_movavg(HE,t,winsize,overlap,0);
        else
            env(vox,:) = HE;
            t_avg = t;
        end    
    end
    ft_progress('close')
    
    t=t_avg;
    
    clear recon
    env(:,isnan(env(1,:))) = [];
    
    disp(['Saving envelope data for ' filestr ' to ' filenames.envelope{subnum}])
    save(filenames.envelope{subnum},'env','t','-v7.3')

    clear dat
    
  end
  
end



%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                         %
%                    P R E P A R E   H M M   D A T A                      %
%                                                                         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if todo.concat || (todo.infer && ~exist(filenames.concat,'file'))
  
  % Load subjects and concatenate:
  env_concat = [];
  subj_inds = [];
  C = 0;
  for subnum = 1:length(data_files)
    load(filenames.envelope{subnum});
    if norm_subs
      env = demean(env,2)./std(env(:));
    end
    env_concat = [env_concat,env];
    subj_inds = [subj_inds,subnum*ones(1,size(env,2))];
    C = C + env*permute(env,[2,1]);
  end
  C = C ./ (length(subj_inds)-1);
  clear env
  
  
  % PCA + whitening - below is equivalent to fastica whitening code but much faster
  [allsvd,MixingMatrix] = eigdec(C,pcadim);
  if whiten
    MixingMatrix = diag(1./sqrt(allsvd)) * MixingMatrix';
  end
  hmmdata =  (MixingMatrix * env_concat)';
  fsample = round(1/mode(diff(t)));
  
  disp(['Saving concatenated envelope data to ' filenames.concat])
  save(filenames.concat,'hmmdata','MixingMatrix','fsample','subj_inds')

end



%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                         %
%                            I N F E R   H M M                            %
%                                                                         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if todo.infer

if ~exist('hmmdata','var')
  load(filenames.concat)
end
 

hmm = osl_hmm_infer(hmmdata,struct('K',nstates,'order',0,'Ninits',nreps,'Hz',fsample));
hmm.MixingMatrix = MixingMatrix;
hmm.fsample = fsample;

% Save results
disp(['Saving inferred HMM to ' filenames.hmm])
save(filenames.hmm,'hmm','subj_inds')

HMMresults = filenames.hmm;

end



%% %%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%
%                                                                         %
%                       O U T P U T   R E S U L T S                       %
%                                                                         %
%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%%

if todo.output

    if ~exist('hmm','var')
        HMMresults = filenames.hmm;
        load(HMMresults)
    end

    statemaps = [filenames.output,'_',output_method];
    switch output_method
        case {'pcorr','var'}
            % Compute partial correlation/variance maps
            stat = zeros(size(hmm.MixingMatrix,2),hmm.K);
            try
                dat = load(data_files{1});            
                dat = dat.oat_stage_results;
            catch
                error([output_method ' currently only supported for OAT results'])
            end
            
            for subnum = unique(subj_inds)
                [~,filestr] = fileparts(data_files{subnum});
                disp(['Computing ' output_method ' maps for ' filestr]);
                load(filenames.envelope{subnum});
                hmm_sub = hmm; hmm_sub.statepath = hmm.statepath(subj_inds==subnum);
                stat = stat + osl_hmm_statemaps(hmm_sub,env,0,output_method);
            end
            stat = stat ./ length(unique(subj_inds));
            disp(['Saving state spatial maps to ' statemaps])
            statemaps = nii_quicksave(stat,statemaps,dat.gridstep,2);

        case 'cov'
            cov_mats = osl_hmm_statemaps(hmm,[],[],'cov');                        
            statemaps = [statemaps '.mat'];
            disp(['Saving state covariance matrices to ' statemaps])
            save(statemaps,'cov_mats')

        otherwise
            warning([output_method ' is not a supported output method']);
    end

end


